package controller

import (
	"context"
	"github.com/gin-gonic/gin"
	"github.com/ollama/ollama/api"
	"github.com/tmc/langchaingo/llms"
	"log"
	"net/http"
	"ollama-qwen/model"
	"ollama-qwen/util"
)

type Chat struct{}

func (ch *Chat) DoChat(ctx *gin.Context) {
	var body model.Chat

	if err := ctx.ShouldBindJSON(&body); err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{
			"msg": err.Error(),
		})
	}

	prompt := util.CreatePrompt()

	data := map[string]any{
		"text": body.Text,
	}

	msg, _ := prompt.FormatMessages(data)

	content := []llms.MessageContent{
		llms.TextParts(msg[0].GetType(), msg[0].GetContent()),
		llms.TextParts(msg[1].GetType(), msg[1].GetContent()),
	}

	llm := util.CreateOllama(ctx, "qwen")

	resp, err := llm.GenerateContent(context.Background(), content)

	if err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{
			"msg": err.Error(),
		})
	}
	ctx.JSON(http.StatusOK, gin.H{
		"code": 0,
		"msg":  "ok",
		"data": resp.Choices[0].Content,
	})
}

func (ch Chat) DoChatStream(ctx *gin.Context) {
	SetSseHeader(ctx)
	var body model.Chat

	if err := ctx.ShouldBindJSON(&body); err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{
			"msg": err.Error(),
		})
	}
	client, err := api.ClientFromEnvironment()
	if err != nil {
		log.Fatal(err)
		return
	}
	req := &api.GenerateRequest{
		Model:  "qwen",
		Prompt: body.Text,
	}
	c := context.Background()
	respFunc := func(resp api.GenerateResponse) error {

		if resp.Done {
			ctx.Request.Context().Done()
		} else {
			ctx.Writer.WriteString(resp.Response)
			ctx.Writer.Flush()
		}
		return nil
	}
	err = client.Generate(c, req, respFunc)
	if err != nil {
		log.Fatal(err)
	}

}

func SetSseHeader(ctx *gin.Context) {
	ctx.Writer.Header().Set("Content-Type", "text/event-stream")
	ctx.Writer.Header().Set("Cache-Control", "no-cache")
	ctx.Writer.Header().Set("Connection", "keep-alive")
}
